import asyncio
import json
import os
import re
import subprocess
import time
from contextlib import asynccontextmanager
from random import random
import random
import docker
import httpx
import jmespath
import uvicorn
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from fastapi import FastAPI, Request, Path
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import HTMLResponse, FileResponse
import ruamel.yaml
from typing import Any

app = FastAPI()
# 允许跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 更改为您的前端页面地址
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

listen_port = 8191
web_port = 8091
data_containers_file_name = "data_containers_list.yml"
data_containers_user_relations_file_name = "data_containers_user_relations_list.yml"
data_containers_use_history_file_name = "data_containers_use_history.yml"
CONTAINER_MAX_NUMBER = 14
# 创建 Docker 客户端对象
client = docker.from_env()
# 创建调度器
scheduler = AsyncIOScheduler()


async def fetch_baidu_api():
    async with httpx.AsyncClient() as client:
        response = await client.get("http://192.168.114.37:7777/check")
        print(f"{time.strftime('%Y-%m-%d %H:%M:%S')}", response.text)
        print(
            f"{time.strftime('%Y-%m-%d %H:%M:%S')}Fetched {len(response.text)} characters from http://192.168.114.37:7777/check")


@app.on_event("startup")
async def startup_event():
    # 添加定时任务，每隔1分钟运行一次
    scheduler.add_job(fetch_baidu_api, "interval", minutes=2)
    scheduler.start()
    print("Scheduler started")


@app.on_event("shutdown")
async def shutdown_event():
    # 停止调度器
    scheduler.shutdown()
    print("Scheduler stopped")


def read_yml(file_path, node=None):
    with open(file_path, "r", encoding="utf-8") as f:
        yaml = ruamel.yaml.YAML(typ='unsafe', pure=True)
        yaml.allow_duplicate_keys = True
        data = yaml.load(f)
    if node:
        return data[node]
    else:
        return data if data else []


def write_to_yml(file_path, data):
    # 顺序写入新的yml文件中,保证格式不混乱
    with open(file_path, "w", encoding="utf-8") as f:
        ruamel.yaml.round_trip_dump(data, f, default_flow_style=False, allow_unicode=True)


def check_port_in_use(port, data_containers_list, data_containers_user_relations_list):
    # 定义要执行的命令
    command = f"netstat -an | grep {port}"
    try:
        # 使用subprocess模块执行命令，并获取输出
        output = subprocess.check_output(command, shell=True, encoding='utf-8')
        # 如果命令输出不为空，则表示端口正在使用
        data_list = []
        for i in output.splitlines():
            i = re.sub(r'\s+', ',', i).split(",")
            data_list.append(i)
        print("拿到的端口使用信息", json.dumps(data_list, ensure_ascii=False))
        status_list = jmespath.search("@[*][-1]", data_list)
        print(status_list)
        if "ESTABLISHED" not in status_list:
            # 没有用户连接了需要释放资源
            search_container = jmespath.search(f"@[?@[1]==`{port}`][]", data_containers_list)
            if search_container:
                data_containers_list.remove(search_container)
                search_container[-1] = False
                data_containers_list.append(search_container)
                try:
                    container = client.containers.get(search_container[0])
                    # 停止容器
                    container.start()
                    print(f"Container {search_container[0]} restart Success.")
                    print(data_containers_list)
                    data_line = jmespath.search(f"@[?@[1]==`{search_container[0]}`][]",
                                                data_containers_user_relations_list)
                    if data_line:
                        data_containers_user_relations_list.remove(data_line)
                    return True
                except docker.errors.NotFound as e:
                    print(f"Container {search_container[0]} not found: {e}")
                except docker.errors.APIError as e:
                    print(f"Error stopping or removing container {search_container[0]}: {e}")
        else:
            return False
    except subprocess.CalledProcessError:
        # 如果命令执行出错，则表示端口没有被使用
        return False

    # 如果命令输出为空，则表示端口没有被使用
    return False

def check_port_use_info(port):
    # 定义要执行的命令
    command = f"netstat -an | grep {port}"
    try:
        # 使用subprocess模块执行命令，并获取输出
        output = subprocess.check_output(command, shell=True, encoding='utf-8')
        # 如果命令输出不为空，则表示端口正在使用
        print(output)
        if output:
            return True
    except subprocess.CalledProcessError:
        # 如果命令执行出错，则表示端口没有被使用
        return False

    # 如果命令输出为空，则表示端口没有被使用
    return False


def stop_and_restart_docker_container(container_id, filter_express=None):
    try:
        # 停止容器
        container = client.containers.get(container_id)
        container.start()
        print(f"Container {container_id} stopped successfully.")
        # 复制文件到容器内指定目录
        with open("mitmproxy_plugin_temp.py", 'rb') as f:
            data = f.read().decode('utf-8')
            data = data.replace("{{filter_express}}", filter_express)
            # print(data)
        with open("mitmproxy_plugin.py", 'wb') as f:
            f.write(data.encode('utf-8'))
        # 重启容器
        # 构建 docker cp 命令
        docker_cp_command = f"docker cp mitmproxy_plugin.py {container_id}:/home/mitmproxy"
        # 使用 os.system 执行命令
        os.system(docker_cp_command)
        print(f"File 'mitmproxy_plugin_temp.py' copied to container '{container_id}' successfully.")
        container.start()
        print("File copied successfully and container restarted.")
        return True
    except docker.errors.NotFound as e:
        print(f"Container {container_id} not found: {e}")
        return False
    except docker.errors.APIError as e:
        print(f"Error stopping or removing container {container_id}: {e}")
        return False


def start_mitmproxy_container(listen_port: int = 8191, web_port: int = 8091) -> tuple[Any, Any]:
    # 容器配置
    container_config = {
        'image': 'ptesting_mitmproxy',  # 镜像名称
        'command': [
            'mitmweb',
            '-s', 'mitmproxy_plugin.py',
            '--set', 'ssl_insecure=true',
            '--listen-host', '192.168.114.37',
            '--listen-port', f'{listen_port}',
            '--web-host', '192.168.114.37',
            '--web-port', f'{web_port}',
            '--set', 'client_certs=client2.pem'
        ],
        "network_mode": "host",
        'detach': True,
        'tty': True
    }

    # 启动容器
    container = client.containers.run(**container_config)

    container = client.containers.get(container.id)
    # 返回容器 ID
    # 5d688e6cba9e
    return container.id[:12]


def business_handle(data):
    """
        data_containers_list：存储容器的基本信息
        data_containers_user_relations_list：存储用户与容器的关系信息
        用户发起请求后，拿到用户的id和用户搜索的关键词
        系统初始化的时候启动5个容器，后面就不再启动，是否被使用，默认是False
        用户访问时随机分配一个容器实例，并将使用状态设置为true
    :param data:
    :return:
    """

    global listen_port, web_port
    search_kw = data["search_kw"]
    user_uuid = data["user_uuid"]
    print("我拿到的search_kw", search_kw)
    print("我拿到的search_kw", user_uuid)
    data_containers_list = read_yml(data_containers_file_name)
    max_listen_port = jmespath.search("max(@[*][1])", data_containers_list)
    max_web_port = jmespath.search("max(@[*][2])", data_containers_list)
    # 扩容的逻辑------------------------
    if max_listen_port:
        listen_port = max_listen_port+1
    if max_web_port:
        web_port = max_web_port+1

    while len(data_containers_list)<CONTAINER_MAX_NUMBER:
        if check_port_use_info(listen_port) or check_port_use_info(web_port):
            listen_port += 1
            web_port += 1
            continue
        new_container_id = start_mitmproxy_container(listen_port, web_port)
        data_containers_list.append([new_container_id, listen_port, web_port, False])
        listen_port += 1
        web_port += 1
    write_to_yml(data_containers_file_name, data_containers_list)
    # 扩容的逻辑------------------------
    data_containers_user_relations_list = read_yml(data_containers_user_relations_file_name)
    print("我拿到的", len(data_containers_user_relations_list))
    print("我拿到的", data_containers_user_relations_list)
    # 启动容器的逻辑
    # 判断是否已经分配过容器
    # 根据user_uuid查询data_containers_user_relations_list，如果有对应的关系数据，则继续使用之前的容器和对应的端口
    # 如果没有对应的关系数据，则需要重新分配容器和端口，在重新分配容器和端口时，需要查询用户使用历史记录表data_containers_use_history.yml
    # 如果用户user_uuid使用历史记录表data_containers_use_history.yml中有使用记录，判断记录中使用的容器是否被使用，如果被使用，则需要重新分配容器和端口，否则继续使用之前的容器和对应的端口
    data_containers_use_history = read_yml(data_containers_use_history_file_name)
    data_containers_use_history = {} if not data_containers_use_history else data_containers_use_history
    history_container_id = data_containers_use_history.get(user_uuid) if data_containers_use_history else ''
    container_id = jmespath.search(f'@[?@[0]==`{user_uuid}`][1]|[0]', data_containers_user_relations_list)
    print("我已经有资源了，直接用之前的资源", container_id)
    if container_id:
        container_detail = jmespath.search(f'@[?@[0]==`{container_id}`][]', data_containers_list)
        print("我已经有资源了，我拿到的", container_detail)
        data_containers_use_history[user_uuid] = container_id
        print(data_containers_use_history)
        write_to_yml(data_containers_use_history_file_name, data_containers_use_history)
        if container_detail[-1] is False:
            stop_and_restart_docker_container(container_id, search_kw)
            index = data_containers_list.index(container_detail)
            data_containers_list.remove(container_detail)
            container_detail[-1] = True
            data_containers_list.insert(index, container_detail)
            write_to_yml(data_containers_file_name, data_containers_list)
            return container_detail
        else:
            stop_and_restart_docker_container(container_id, search_kw)
            return container_detail
    else:
        use_container_list = jmespath.search('@[?@[-1]==`false`]', data_containers_list)
        if use_container_list:
            if history_container_id:
                history_use_container = jmespath.search(f'@[?@[0]==`{history_container_id}`]|[0]', use_container_list)
                container_detail = history_use_container if history_use_container else random.choice(use_container_list)
            else:
                container_detail = random.choice(use_container_list)
            print("新分配的容器", container_detail)
            print("新分配的容器", container_detail[0])
            data_containers_use_history[user_uuid] = container_detail[0]
            stop_and_restart_docker_container(container_detail[0], search_kw)
            index = data_containers_list.index(container_detail)
            data_containers_list.remove(container_detail)
            container_detail[-1] = True
            data_containers_list.insert(index, container_detail)
            if [user_uuid, container_detail[0]] not in data_containers_user_relations_list:
                data_containers_user_relations_list.append([user_uuid, container_detail[0]])
            write_to_yml(data_containers_file_name, data_containers_list)
            write_to_yml(data_containers_user_relations_file_name, data_containers_user_relations_list)
            write_to_yml(data_containers_use_history_file_name, data_containers_use_history)
            print("数据写入完成", data_containers_use_history)
            return container_detail
        else:
            return []


@app.get("/")
async def get_index():
    with open("index.html", "r", encoding="utf-8") as file:
        html_content = file.read()
    return HTMLResponse(content=html_content, status_code=200)


@app.post("/get_content")
async def get_content(request: Request):
    data = await request.json()
    print(data)
    client_ip = request.client.host
    print(client_ip)
    container_detail = business_handle(data)
    # 在这里处理接收到的输入文本，并生成需要返回的内容
    # 这里只是简单地将输入文本返回，您可以根据实际需求处理数据并返回相应的内容
    if container_detail:
        return {
            "web_url": f"http://192.168.114.37:{container_detail[2]}"
        }
    else:
        return {
            "content": "<h3>错误提示</h3>暂无可以使用的资源，请耐心等待"
        }


@app.get("/check")
async def check_content():
    """
        检查端口的使用情况
    :return:
    """
    data_containers_list = read_yml(data_containers_file_name)
    data_containers_user_relations_list = read_yml(data_containers_user_relations_file_name)
    port_str = ""
    port_list = jmespath.search("@[*][1]", data_containers_list)
    for port in port_list:
        result = check_port_in_use(port, data_containers_list, data_containers_user_relations_list)
        if result:
            write_to_yml(data_containers_file_name, data_containers_list)
            write_to_yml(data_containers_user_relations_file_name, data_containers_user_relations_list)
            port_str += f"{port},"
    if port_str:
        return {"content": f"{port_str}对应的容器已释放"}
    else:
        return {"content": f"没有可以释放资源的容器"}


@app.get("/favicon.ico")
def get_favicon():
    file_path = "favicon.ico"  # 替换成你的 favicon.ico 文件路径

    return FileResponse(file_path, media_type="image/x-icon")


@app.get("/download/{filename}")
def download_file(filename: str = Path(...)):
    file_path = f".mitmproxy/{filename}"  # 替换为你的文件所在的文件夹路径

    return FileResponse(file_path, filename=filename)


if __name__ == '__main__':
    uvicorn.run(app, port=7777, host="192.168.114.37")
    # nohup python3 mitmweb_start.py > output.log 2>&1 &
